import torch
from torch import nn
from torch.nn import functional as F
import dltools

batchSize = 32
num_steps = 35
train_iter, vocab = dltools.load_data_time_machine(batch_size=batchSize, num_steps=num_steps)
num_hiddens = 256
num_layers = 2
# bidirectional=True 双向神经网络主要用于特征提取
gru_layer = nn.GRU(len(vocab),num_hiddens, num_layers=num_layers)
model = dltools.RNNModel(gru_layer, len(vocab))
model = model.to(dltools.try_gpu())
num_epochs, lr=1000, 0.8
dltools.train_ch8(model, train_iter, vocab, lr, num_epochs, dltools.try_gpu())